﻿from math import log2
import operator



def get_entropy(labels, indices):
    """计算列表的平均信息量

    Parameters
    ----------
    labels + indices: [] + []
        需要统计的类别

    Returns
    -------
    entropy:
        平均信息量
    """

    counts = {}
    for l in [labels[i] for i in indices]:
        counts[l] = counts.get(l, 0) + 1;

    entropy = 0.0
    for i in counts:
        prob = float(counts[i]) / len(indices)
        entropy -= prob * log2(prob)
    return entropy


def split_rows_by_col(props, rows, col):
    """根据选择的列分割行

    Parameters
    ----------
    props: [[]]
        属性
    rows:
        需要分割数据的行号
    col:
        按照第col列的值进行分割

    returns
    -------
    splited: {prop: rows}
        按照属性划分好的行
    """

    splited = {}
    for i in rows:
        splited.setdefault(props[i][col], []).append(i)
    return splited


def choose_col_to_split(data, rows, cols, labels):
    '''选择一个列（属性）作为划分标准

    Parameters
    ----------
    props + rows + cols 
        组成属性集
    labels + rows
        组成标签集

    Returns
    -------
    best_col
        被选中的列
    best_splited
        根据选择的列划分
    '''

    (best_entropy, best_col, best_splited) = (get_entropy(labels, rows), -1, {})

    for i in range(len(cols)):

        splited = split_rows_by_col(data, rows, cols[i])

        entropy = 0.0
        for key in splited:
            prob = len(splited[key]) / float(len(rows))
            entropy += prob * get_entropy(labels, splited[key])

        if best_col == -1 or entropy < best_entropy:
            best_entropy = entropy
            best_col = cols[i]
            best_splited = splited

    return (best_col, best_splited)
  

def count_labels(labels, indices):
    '''统计labels[indices]的数目
    '''
    count = {}
    for l in [labels[i] for i in indices]:
        count[l] = count.get(l, 0) + 1
    sorted_count = sorted(count.items(), key=operator.itemgetter(1), reverse=True)
    return sorted_count


def _createTree(props, rows, cols, labels):
    
    sorted_count = count_labels(labels, rows)
    print(sorted_count)
    print(len(sorted_count))
    if len(cols) >= 1 and len(sorted_count) > 1:
        (col, spilted) = choose_col_to_split(props, rows, cols, labels)
        tree = [col, {}]
        subcols = cols.copy()
        subcols.remove(col)
        for key in spilted.keys():
            tree[1][key] = _createTree(props, spilted[key], subcols, labels)
        
        return tree    
    else:
        return sorted_count[0][0]
    
        

def create(data, labels):
    '''创建决策树
    Parameters
    ----------
    data: [[]]
        属性数据
    labels: []
        属性对应的类型

    Returns
    -------
    tree: [col, {prop: tree/label}]
        决策树
    '''
    return _createTree(data, [i for i in range(len(data))],
        [i for i in range(len(data[0]))], labels)               
    
def _classify(x, tree):    
    subtree = tree[1][x[tree[0]]]
    if isinstance(subtree[1], dict):
        return classify(x, subtree);
    else:
        return subtree





class DecisionTree:
    def __init__(self):
        pass

    def train(self, xs, cs):
        pass

    def classify(self, xs):
        cs = []
        for x in xs:
            cs.append(_classify(x, self.tree))
        return cs

